import base64

import cv2
import ray
import requests
from fastapi import FastAPI, UploadFile, File
from ray import serve

from rapidocr_onnxruntime import RapidOCR, VisRes

app = FastAPI()

@serve.deployment
@serve.ingress(app)
class RapidOCRDeployment:

    def __init__(self):
        self.engine = RapidOCR()
        self.vis = VisRes()
        self.font_path = "resources/fonts/FZYTK.TTF"
    @app.get("/")
    def root(self):
        return "Welcome to RapidOCR!"

    @app.post("/predict")
    async def predict(self, file: UploadFile = File(...)):
        contents = await file.read()
        result, elapse_list = self.engine(contents)
        boxes, txts, scores = list(zip(*result))
        print(txts)
        res = self.vis(contents, boxes, txts, scores, self.font_path)
        # 将图像转换为 Base64 编码
        retval, buffer = cv2.imencode('.jpg', res)
        image_bytes = buffer.tobytes()
        base64_image = base64.b64encode(image_bytes).decode("utf-8")
        return {"result":base64_image}

rapidocr_app = RapidOCRDeployment.bind()
# serve.run(RapidOCRDeployment.bind(), route_prefix="/rapidocr", port=13036)

# image_path = "tests/images/1.png"
# files = {"file": open(image_path, "rb")}
# response = requests.post(url = "http://localhost:13036/rapidocr/predict" , files=files)
# print(response.json())